import torch 
from preprocessing.load_dicom import load_dicom
import SimpleITK as sitk 
# from resample.resample import resample_itk_list
from preprocessing.resample_standard import resample_standard_list
import numpy as np 
import torch.nn as nn 

from monai.inferers.inferer import SlidingWindowInferer

#need_series = ["t2_tirm_tra_dark-fluid", "t1_fl2d_tra", "adc"]
resample_size = (24, 256, 256)

def get_t1_series(itk_list):
    for d in itk_list:
        d[0] = d[0].strip()
        # if d[0] == "t1_fl2d_tra":
        if 't1' in d[0] and 'tra' in d[0]:
            return d[1]
    
    return None 

def get_resampled_target_data(data):
    raw_data = [None, None, None]
    t1_shape = None 
    
    # for d in data:
    #     d[0] = d[0].strip()
    #     if d[0] == "t2_tirm_tra_dark-fluid":
    #         raw_data[1] = d[1]
    #     elif d[0] == "t1_fl2d_tra":
    #         raw_data[0] = d[1]
    #     elif "ADC" in d[0]:
    #         raw_data[2] = d[1]
      
    #     if d[0] == need_series[1]:
    #         t1_arr = sitk.GetArrayFromImage(d[1])
    #         t1_shape = t1_arr.shape
    
    # 放宽筛选条件
    for d in data:
        d[0] = d[0].strip()
        
        # 下面这行修改了
        if 't2' in d[0] and 'tra_dark-fluid' in d[0]:
            raw_data[1] = d[1]
        # 下面这行修改了
        elif 't1' in d[0] and 'tra' in d[0]:
            raw_data[0] = d[1]
        elif "ADC" in d[0]:
            raw_data[2] = d[1]
            
        if 't1'in d[0] and 'tra' in d[0]:
            t1_arr = sitk.GetArrayFromImage(d[1])
            t1_shape = t1_arr.shape
            
    assert t1_shape is not None 

    # resample operation
    # resampled_data = resample_itk_list(raw_data, out_size=resample_size)
    resampled_data = resample_standard_list(raw_data)
    
    # converted the resampled data to the numpy array
    data_arr = [sitk.GetArrayFromImage(d) for d in resampled_data]
    data_arr = np.stack(data_arr, axis=0)

    return data_arr, t1_shape

def segment(image, model, raw_size):
    
    window_infer = SlidingWindowInferer(roi_size=(24, 192, 192),
                                                sw_batch_size=2,
                                                overlap=0.5,
                                                mode="gaussian",
                                                progress=True)
    
    pred = window_infer(image, model)
    d, w, h = raw_size

    pred = torch.nn.functional.interpolate(pred, mode="trilinear", size=(d, w, h))

    # pred = model(image).argmax(dim=1, keepdim=True).float()

    pred = pred.argmax(dim=1)

    # import matplotlib.pyplot as plt 

    # for i in range(20):
    #     plt.subplot(1, 2, 1)
    #     plt.imshow(image[0, 0, i], cmap="gray")
    #     plt.subplot(1, 2, 2)
    #     plt.imshow(pred[0, i], cmap="gray")
    #     plt.show()

    pred = pred[0]

    return pred 


def norm(data):

    data = data.astype(np.float32)
    data = torch.from_numpy(data).float()[None, ]

    data = (data - data.mean()) / data.std()
    return data 


def norm_v2(data):

    data = data.astype(np.float32)
    data = torch.from_numpy(data).float()
    for c in range(data.shape[0]):
        datac = data[c]
        data[c] = (datac - datac.mean()) / datac.std()
    
    data = data[None, ]

    return data

def run_segmentation(itk_list, model):
    input_data, raw_shape = get_resampled_target_data(itk_list)

    input_data = input_data[:2] ## only get the t1 and flair modality
    input_data = norm_v2(input_data)

    print(f"input data shape is {input_data.shape}")
    print("*" * 100)
    # The shape of t1 is {d, w, h}, and it has beed resampled to the same shape with t1 modality.
    pred = segment(input_data, model, raw_shape)

    return pred 


